package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"
)

func main() {
	flag.Parse()
	roots := flag.Args()
	if len(roots) == 0 {
		roots = []string{"."}
	}

	fileSizes := make(chan int64)
	var n sync.WaitGroup
	for _, root := range roots {
		n.Add(1)
		go walkDir(root, &n, fileSizes)
	}
	go func() {
		n.Wait()
		close(fileSizes)
	}()

	// Print the results.
	var nfiles, nbytes int64
	for size := range fileSizes {
		nfiles++
		nbytes += size
	}
	printDiskUsage(nfiles, nbytes)
}

func printDiskUsage(nfiles, nbytes int64) {
	fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
}

var done = make(chan struct{})
var sema = make(chan struct{}, 20)

func canceled() bool {
	select {
	case <-done:
		return true
	default:
		return false
	}
}

func walkDir(dir string, n *sync.WaitGroup, filesSizes chan<- int64) {
	defer n.Done()
	if canceled() {
		return
	}
	for _, file := range dirents(dir) {
		if file.IsDir() {
			n.Add(1)
			subdir := filepath.Join(dir, file.Name())
			go walkDir(subdir, n, filesSizes)
		} else {
			filesSizes <- file.Size()
		}

	}
}

func dirents(dir string) []os.FileInfo {
	select {
	case sema <- struct{}{}:
	case <-done:
		return nil
	}
	defer func() {
		<-sema
	}()
	readDir, err := ioutil.ReadDir(dir)
	if err != nil {
		return nil
	}
	return readDir

}
